import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

'''
        任务：鸢尾花识别
'''

DATA_FILE = './data/Iris.csv'

SPECIES_LABEL_DICT = {
    'Iris-setosa': 0, # 山鸢尾
    'Iris-versicolor': 1, # 变色鸢尾
    'Iris-virginica': 2 # 维吉尼亚鸢尾
}

# 使用的特征列
FEAT_COLS = ['SepalLengthCm', 'SepalWidthCm', 'PetalLengthCm', 'PetalWidthCm']

def main():
    # 1. 读取数据
    iris_data = pd.read_csv(DATA_FILE, index_col='Id')
    iris_data['Label'] = iris_data['Species'].map(SPECIES_LABEL_DICT)

    # 获取数据集特征
    X = iris_data[FEAT_COLS].values
    # 获取数据集标签
    y = iris_data['Label'].values
    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=1/3, random_state=10)
    # 声明变量
    knn_model = KNeighborsClassifier(n_neighbors=3)

    # 2. 训练模型
    knn_model.fit(X_train, y_train)
    # 3. 模型评估
    score = knn_model.score(X_test, y_test)
    print('模型准确率：{:.2f}%'.format(score * 100))


if __name__ == '__main__':
    main()